package optimization;

import Jama.Matrix;
import static optimization.LinearAlgebra.*;

public abstract class NewtonNonlinearSolver {

	public static double[] newtonNonlinearSolver(VectorFunction f, double[] x0, double[] y)
	{
		return newtonNonlinearSolver(f, x0, y, 0.001);
	}
	/**
	 * Solve Ax=B or J(x)*dx = -r(x) 
	 * @param jacob
	 * @param y
	 * @param yt
	 * @return p or dx
	 */
	public static double[] solveJPR(double[] jacob, double[] y, double[] yt)
	{
		int n = y.length;
		int m = jacob.length / n;
		
		double[] r = vSub(y, yt); // r = target - current
		Matrix B = new Matrix(r, n);
		
		Matrix A = new Matrix(jacob, m).transpose();

		// Solve Ax=B or J(x)*dx = -r(x) 
		Matrix R = A.lu().solve(B);
		double[] p = R.getRowPackedCopy();
		return p;
	}
	
	/**
	 * Solve nonlinear system of equations using Newton's method
	 * 
	 * !! Currently limits parameter space to [0, 6] x [0, 1]
	 * @param f
	 * @param x0
	 * @param y 
	 * @return
	 */
	public static double[] newtonNonlinearSolver(VectorFunction f, double[] x0, double[] y, double maxerror)
	{
		System.out.println("NEWTON NONLSOLV start " + vPrint(x0));
		double[] x = x0;
		int maxIter = 100;
		
		TabuList tabu = new TabuList(8, maxerror/10);
		
		double[] h = vRepeat(x.length, 0.1);
		int iter = 0;
		while (maxIter-- > 0) {
			iter++;
			
			// 1. Solve J(x)*p = -r(x) for p
			double[] fx = f.evaluate(x);
			double[] r = vSub(y, fx); // r = target - current
			Matrix B = new Matrix(r, r.length);
			
			// B E err TSm TSs
//			System.out.print("NLS(" + iter + ") X: " + vPrint(x) + " Y: " + vPrint(fx) + " R: " + vPrint(r) + " ||R||: " + String.format("%.6f", B.norm2()));
			System.out.print(iter + " " + vPrint(x,false) + " " + vPrint(fx,false) + " " + String.format("%.6f", B.norm2()));
			
			// stop if close enough
			if (B.norm2() < maxerror) break;
			
			// Compute jacobian
			double[] jacob = jacobianForwardDifferences(f, fx, x, h);
//			System.out.println("JACOB " + vPrint(jacob));
			
			Matrix A = new Matrix(jacob, x.length).transpose();

			// Solve Ax=B or J(x)*dx = -r(x) 
			Matrix R = A.lu().solve(B);
			double[] p = R.getRowPackedCopy();
	//		System.out.println("DeltaX " + vPrint(p));
			
			// 2. x' = x + p    Update X
			double alpha = 1;
			if (x[0]+p[0] < 0) alpha = -x[0]/p[0];
			if (x[0]+p[0] > 6) alpha = (6-x[0])/p[0];
			if (x[1]+p[1] < 0) alpha = -x[1]/p[1];
			if (x[1]+p[1] > 1) alpha = (1-x[1])/p[1];
//			if (alpha != 1) { System.out.println("ALPHA " + alpha); }
			double[] nx = vAdd(x, vMul(alpha,p));
			if (tabu.accept(nx)) {
				x = nx;
			} else {
				// not accepted, try halving
				boolean ok = false;
				for (int i = 0; i < 4; i++) {
					alpha = alpha/2;
					nx = vAdd(x, vMul(alpha,p));
					if (tabu.accept(nx)) {
						x = nx;
						ok = true;
						break;
					}
				}
				if (!ok) {
					System.out.println("TABU BLOCKED, STOP.");
					return x;
				}
			}
//			System.out.println("NewX " + vPrint(x));
	
//			System.out.println("NLS("+iter+") JdX=R J: " + vPrint(jacob) + " dX: " + vPrint(p) );
//			System.out.println("  J: " + vPrint(jacob) + " dX: " + vPrint(p) );
			System.out.println(" " + vPrint(jacob,false) + " " + vPrint(p,false) +";");
			// enforce limits on values 
			if (x[0] < 0) x[0] = 0;
			if (x[0] > 6) x[0] = 6;
			if (x[1] < 0) x[1] = 0;
			if (x[1] > 1) x[1] = 1;
			
			// if the values are far away, create new random starting point
			if (x[0] == 0 && x[1] == 0) {
				x[0] = Math.random()*6;
				x[1] = Math.random();
				System.out.println("RANDOM RESTART");
			}
			
			// 3. repeat
		}
		
		System.out.println("NLS-RETURN " + vPrint(x));
		return x;
	}
	
}
